import os
import numpy as np
from PIL import Image, ImageEnhance, ImageOps, ImageFilter
import random
import math
from tqdm import tqdm


class SimpleImageAugmenter:
    def __init__(self, input_dir, output_dir):
        """
        初始化图像增强类

        参数:
        input_dir: 输入图像目录路径
        output_dir: 增强后图像保存目录路径
        """
        self.input_dir = input_dir
        self.output_dir = output_dir

        # 创建输出目录（如果不存在）
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

    def load_images(self):
        """加载输入目录中的所有图像文件"""
        image_files = []
        valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp']

        for filename in os.listdir(self.input_dir):
            ext = os.path.splitext(filename)[1].lower()
            if ext in valid_extensions:
                image_files.append(os.path.join(self.input_dir, filename))

        return image_files

    def rotate_image(self, image, max_angle=30):
        """
        随机旋转图像

        参数:
        image: PIL图像对象
        max_angle: 最大旋转角度
        """
        angle = random.uniform(-max_angle, max_angle)
        return image.rotate(angle, resample=Image.BICUBIC, expand=False)

    def adjust_brightness(self, image, factor_range=(0.7, 1.3)):
        """
        调整图像亮度（模拟曝光变化）

        参数:
        image: PIL图像对象
        factor_range: 亮度调整因子范围，如(0.7, 1.3)表示亮度在原图70%到130%之间
        """
        factor = random.uniform(factor_range[0], factor_range[1])
        enhancer = ImageEnhance.Brightness(image)
        return enhancer.enhance(factor)

    def add_noise(self, image, noise_level=0.05):
        """
        添加随机噪点

        参数:
        image: PIL图像对象
        noise_level: 噪点程度，取值范围0-1
        """
        img_array = np.array(image)

        # 创建高斯噪声
        noise = np.random.normal(0, 25, img_array.shape)
        noisy_img = img_array + noise_level * noise
        noisy_img = np.clip(noisy_img, 0, 255).astype(np.uint8)

        # 随机添加椒盐噪声
        if random.random() > 0.5:
            h, w, _ = img_array.shape
            noise_pixels = int(noise_level * 0.1 * h * w)

            for _ in range(noise_pixels):
                y = random.randint(0, h - 1)
                x = random.randint(0, w - 1)
                # 黑色或白色噪点
                noisy_img[y, x] = [random.choice([0, 255])] * 3

        return Image.fromarray(noisy_img)

    def add_shadow(self, image, shadow_intensity=(0.4, 0.7)):
        """
        添加随机阴影效果

        参数:
        image: PIL图像对象
        shadow_intensity: 阴影强度范围
        """
        width, height = image.size
        img_array = np.array(image)

        # 创建阴影掩码
        mask = np.ones((height, width), dtype=np.float32)

        # 随机选择阴影类型
        shadow_type = random.choice(['gradient', 'ellipse', 'rectangle'])

        if shadow_type == 'gradient':
            # 创建渐变阴影
            direction = random.choice(['top', 'right', 'bottom', 'left'])
            for y in range(height):
                for x in range(width):
                    if direction == 'top':
                        grad = y / height
                    elif direction == 'right':
                        grad = (width - x) / width
                    elif direction == 'bottom':
                        grad = (height - y) / height
                    else:  # left
                        grad = x / width

                    mask[y, x] = 1.0 - (1.0 - grad) * random.uniform(shadow_intensity[0], shadow_intensity[1])

        elif shadow_type == 'ellipse':
            # 创建椭圆阴影
            center_x = random.randint(0, width)
            center_y = random.randint(0, height)
            rx = random.randint(width // 4, width // 2)
            ry = random.randint(height // 4, height // 2)

            for y in range(height):
                for x in range(width):
                    dist = ((x - center_x) / rx) ** 2 + ((y - center_y) / ry) ** 2
                    if dist < 1.0:
                        intensity = random.uniform(shadow_intensity[0], shadow_intensity[1])
                        mask[y, x] = 1.0 - intensity * (1.0 - dist)

        else:  # rectangle
            # 创建矩形阴影
            x1 = random.randint(0, width - 10)
            y1 = random.randint(0, height - 10)
            x2 = random.randint(x1 + 10, width)
            y2 = random.randint(y1 + 10, height)

            intensity = random.uniform(shadow_intensity[0], shadow_intensity[1])
            mask[y1:y2, x1:x2] = 1.0 - intensity

        # 应用阴影
        for c in range(3):  # 对RGB三个通道应用
            img_array[:, :, c] = img_array[:, :, c] * mask

        return Image.fromarray(img_array.astype(np.uint8))

    def horizontal_flip(self, image):
        """水平翻转图像"""
        if random.random() > 0.5:
            return image.transpose(Image.FLIP_LEFT_RIGHT)
        return image

    def adjust_contrast(self, image, factor_range=(0.7, 1.3)):
        """调整对比度"""
        factor = random.uniform(factor_range[0], factor_range[1])
        enhancer = ImageEnhance.Contrast(image)
        return enhancer.enhance(factor)

    def adjust_saturation(self, image, factor_range=(0.7, 1.3)):
        """调整饱和度"""
        factor = random.uniform(factor_range[0], factor_range[1])
        enhancer = ImageEnhance.Color(image)
        return enhancer.enhance(factor)

    def apply_blur(self, image, max_radius=2.0):
        """应用模糊效果"""
        radius = random.uniform(0, max_radius)
        return image.filter(ImageFilter.GaussianBlur(radius=radius))

    def random_crop(self, image, crop_ratio=(0.8, 0.95)):
        """随机裁剪图像"""
        width, height = image.size
        ratio = random.uniform(crop_ratio[0], crop_ratio[1])

        crop_width = int(width * ratio)
        crop_height = int(height * ratio)

        left = random.randint(0, width - crop_width)
        top = random.randint(0, height - crop_height)
        right = left + crop_width
        bottom = top + crop_height

        cropped = image.crop((left, top, right, bottom))
        return cropped.resize((width, height), Image.BICUBIC)

    def augment(self, image_path, num_outputs=5):
        """
        对一张图像进行多种增强处理

        参数:
        image_path: 图像文件路径
        num_outputs: 每张原始图像生成的增强图像数量
        """
        try:
            # 读取原始图像
            image = Image.open(image_path).convert('RGB')
            filename = os.path.splitext(os.path.basename(image_path))[0]

            # 对每张图像创建多个增强版本
            for i in range(num_outputs):
                # 复制原图
                aug_image = image.copy()

                # 定义可能应用的增强操作及其概率
                augmentations = [
                    (self.rotate_image, 0.7),
                    (self.adjust_brightness, 0.7),
                    (self.add_noise, 0.5),
                    (self.add_shadow, 0.5),
                    (self.horizontal_flip, 0.5),
                    (self.adjust_contrast, 0.6),
                    (self.adjust_saturation, 0.6),
                    (self.apply_blur, 0.4),
                    (self.random_crop, 0.4)
                ]

                # 应用随机选择的增强操作
                for aug_func, prob in augmentations:
                    if random.random() < prob:
                        aug_image = aug_func(aug_image)

                # 保存增强后的图像
                output_path = os.path.join(self.output_dir, f"{filename}_aug_{i}.jpg")
                aug_image.save(output_path, "JPEG", quality=95)

            return True

        except Exception as e:
            print(f"处理图像 {image_path} 时出错: {e}")
            return False

    def process_all_images(self, num_outputs=5):
        """
        处理所有输入图像

        参数:
        num_outputs: 每张原始图像生成的增强图像数量
        """
        # 加载所有图像
        image_files = self.load_images()

        if not image_files:
            print(f"在 {self.input_dir} 目录中未找到图像文件")
            return

        print(f"找到 {len(image_files)} 张图像，开始处理...")

        # 处理每张图像
        success_count = 0
        for image_path in tqdm(image_files):
            if self.augment(image_path, num_outputs):
                success_count += 1

        print(f"处理完成! 成功增强 {success_count}/{len(image_files)} 张图像")
        print(f"增强后的图像已保存到 {self.output_dir} 目录")


# 使用示例
if __name__ == "__main__":
    # 配置输入和输出目录
    input_directory = "D:\parking"  # 原始图像目录
    output_directory = r"D:\augmented_images2"  # 增强后图像保存目录

    # 创建增强器
    augmenter = SimpleImageAugmenter(input_directory, output_directory)

    # 每张图像生成20个增强版本
    augmenter.process_all_images(num_outputs=20)